# -*- coding: utf-8 -*-
import uuid
import json
from django.core.management.base import BaseCommand
from django.db import transaction
from django.utils import timezone
from django.db.models import Q
from core.models import EmailMessage, Mailbox, Scene, ExtractionRun, process_extraction_and_merge
from core.views import _render_prompt, _call_ai_provider

class Command(BaseCommand):
    help = "Process pending emails: lock, run extraction via AI, and merge results."

    def add_arguments(self, parser):
        parser.add_argument('--mailbox', type=int, help='Only process a specific Mailbox ID')
        parser.add_argument('--scene', type=int, help='Only use a specific Scene ID for processing')
        parser.add_argument('--batch', type=int, default=10, help='Batch size to claim at once (default 10)')
        parser.add_argument('--timeout', type=int, help='Override AI timeout seconds for provider')

    def handle(self, *args, **options):
        mailbox_id = options.get('mailbox')
        scene_id = options.get('scene')
        batch = max(1, min(int(options.get('batch') or 10), 100))
        timeout = options.get('timeout')

        mailboxes = Mailbox.objects.filter(enabled=True)
        if mailbox_id:
            mailboxes = mailboxes.filter(pk=mailbox_id)
        scenes = Scene.objects.filter(enabled=True)
        if scene_id:
            scenes = scenes.filter(pk=scene_id)
        if not scenes.exists():
            self.stdout.write(self.style.WARNING('No enabled scenes to process.'))
            return

        # 简化：按主题匹配场景（沿用 run_single 中逻辑）
        def select_scene_for(email_obj):
            subj = (email_obj.subject or '').strip()
            for sc in scenes.order_by('priority'):
                try:
                    pattern = sc.email_type or ''
                    if pattern and re.search(pattern, subj, re.I):
                        return sc
                except Exception:
                    continue
            return None

        total_ok = 0
        total_fail = 0
        total_updated = 0
        while True:
            # 领取一批任务
            with transaction.atomic():
                qs = EmailMessage.objects.select_for_update(skip_locked=True).filter(status='pending')
                if mailbox_id:
                    qs = qs.filter(mailbox_id=mailbox_id)
                emails = list(qs.order_by('id')[:batch])
                now = timezone.now()
                token = uuid.uuid4().hex
                for e in emails:
                    e.status = 'running'
                    e.lock_token = token
                    e.locked_at = now
                if emails:
                    EmailMessage.objects.bulk_update(emails, ['status', 'lock_token', 'locked_at'])
            if not emails:
                break

            for e in emails:
                try:
                    sc = select_scene_for(e)
                    if not sc:
                        # 无匹配场景，直接记为成功但不更新宽表
                        e.status = 'done'
                        e.processed_at = timezone.now()
                        e.last_error = ''
                        e.lock_token = ''
                        e.locked_at = None
                        e.save(update_fields=['status', 'processed_at', 'last_error', 'lock_token', 'locked_at'])
                        total_ok += 1
                        continue
                    # 渲染提示词与调用AI
                    prompt = _render_prompt(sc, e)
                    data, reasoning = _call_ai_provider(sc.ai_provider, prompt, timeout=timeout)
                    # 合并与记录
                    result = process_extraction_and_merge(scene=sc, mailbox=e.mailbox, email=e, results=data, reasoning=reasoning)
                    e.status = 'done'
                    e.processed_at = timezone.now()
                    e.last_error = ''
                    e.lock_token = ''
                    e.locked_at = None
                    e.save(update_fields=['status', 'processed_at', 'last_error', 'lock_token', 'locked_at'])
                    total_ok += 1
                    if result.get('updated'):
                        total_updated += 1
                except Exception as ex:
                    e.status = 'failed'
                    e.last_error = str(ex)[:1000]
                    e.lock_token = ''
                    e.locked_at = None
                    e.processed_at = timezone.now()
                    e.save(update_fields=['status', 'last_error', 'lock_token', 'locked_at', 'processed_at'])
                    total_fail += 1

        self.stdout.write(self.style.SUCCESS(f"Done. ok={total_ok} fail={total_fail} updated={total_updated}"))